import yappi, re

RE_FUNC = re.compile('.*?/(?P<func>[^/]+?/[^/]+?:\d+?)')

class Profiler:


    def __init__(self, habilitadas=None):
        self.habilitadas = ['workspace', 'pygame', 'threading']
        if habilitadas:
            self.habilitadas += habilitadas
        self.result = ','.join(['name',
                                'count',
                                'total',
                                'totalsub',
                                'tavg',
                                'tsubavg',
                                'diffavg']) + '\n'

    def ok(self, s):
        return any([(x in s) for x in self.habilitadas])



    def _process_stat(self, stat):
        stat = list(stat)
        func_name = stat[0]
        if self.ok(func_name):
            match_obj = RE_FUNC.match(func_name)
            if match_obj:
                stat[0] = match_obj.groupdict()['func']
            #promedios
            stat.append(stat[2] / stat[1])
            stat.append(stat[3] / stat[1])
            #diferencia
            stat.append(stat[4] - stat[5])
            stat = [ "%0.5f" % x if type(x) == float else x for x in stat]
            stat = ','.join([str(x) for x in stat])
            self.result += stat + '\n'
        else:
            print func_name

    def profile(self, main, filename='res.csv'):
        """ Receive function to profile, writes result to filename(res.csv) """
        yappi.start(False)
        main()
        yappi.stop()
        yappi.enum_stats(self._process_stat)
        fileout = file(filename, 'w')
        fileout.write(self.result)
        fileout.close()
